Description

This PR refactors the training loop from a script to a class for better organization and reusability. The training loop code was originally located in the src/main.py file and has been moved into a new class named MNISTTrainer. The src/api.py file has also been updated to import and use the new class for prediction.

Summary of Changes

• Created a new class named MNISTTrainer in src/main.py to encapsulate the training loop functionality.
• Moved the code for loading and preprocessing the MNIST dataset into a method named load_data in the MNISTTrainer class.
• Moved the code for defining the PyTorch model into a method named define_model in the MNISTTrainer class.
• Updated the Net class in src/main.py to accept a trainloader parameter in its constructor and use it for training the model.
• Updated the script at the bottom of src/main.py to create an instance of MNISTTrainer, call its methods to load the data and define the model, and then pass the loaded data to the Net class for training.
• Updated the import statement in src/api.py to import the MNISTTrainer class instead of the Net class.
• Updated the code for loading the model in src/api.py to create an instance of MNISTTrainer, call its methods to load the data and define the model, and then load the model's state from the saved file.
• Updated the documentation in README.md to reflect the changes in the codebase, specifically explaining the new MNISTTrainer class and how it is used in src/main.py and src/api.py.

Fixes #6.

🎉 Latest improvements to Sweep:

• Sweep can now passively improve your repository! Check out Rules to learn more.

💡 To get Sweep to edit this pull request, you can:

• Comment below, and Sweep can edit the entire PR
• Comment on a file, Sweep will only modify the commented file
• Edit the original issue to get Sweep to recreate the PR from scratch